#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/core/core.hpp>
#include <math.h>
#include <iostream>
#include <opencv2/opencv.hpp>

#define armor_height 135
#define armor_width 53

using namespace std;
using namespace cv;

int low_h = 10;
int low_s = 120;
int low_v = 245;

int max_h = 25;
int max_s = 255;
int max_v = 255;

double data_m[9] = {1.2853517927598091e+03, 0., 3.1944768628958542e+02, 0.,
                    1.2792339468697937e+03, 2.3929354061292258e+02, 0., 0., 1.};
double data_d[5] = {-6.3687295852461456e-01, -1.9748008790347320e+00,
                    3.0970703651800782e-02, 2.1944646842516919e-03, 0.};
Mat camera_matrix = Mat(3, 3, CV_64FC1, data_m);
Mat distortion_coefficients = Mat(5, 1, CV_64FC1, data_d);

vector<Point3f> armor = vector<Point3f>{
    Point3f(-armor_height / 2, -armor_width / 2, 0),
    Point3f(armor_height / 2, -armor_width / 2, 0),
    Point3f(armor_height / 2, armor_width / 2, 0),
    Point3f(-armor_height / 2, armor_width / 2, 0)};
vector<Point2f> points;

Mat r_vec = Mat::zeros(3, 1, CV_64FC1);
Mat t_vec = Mat::zeros(3, 1, CV_64FC1);
Mat dis = Mat::ones(1, 1, CV_64FC1);

void draw(Mat src, Mat dst)
{
    vector<RotatedRect> rr;
    vector<RotatedRect> vRec;
    vector<vector<Point>> light_contours;
    vector<Vec4i> hierarchy;
    RotatedRect Rec;
    findContours(src, light_contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    for (int i = 0; i < light_contours.size(); i++)
    {
        float contours_area = contourArea(light_contours[i]);
        if (contours_area < 15 || light_contours[i].size() <= 10)
            continue;

        RotatedRect Rec = fitEllipse(light_contours[i]);
        rr.push_back(Rec);
    }

    for (size_t i = 0; i < rr.size(); i++)
    {
        for (size_t j = i + 1; (j < rr.size()); j++)
        {
            float contour_angle = abs(rr[i].angle - rr[j].angle);
            if (contour_angle >= 10)
                continue;

            float contour_len1 = abs(rr[i].size.height - rr[j].size.height) / max(rr[i].size.height, rr[j].size.height);
            float contour_len2 = abs(rr[i].size.width - rr[j].size.width) / max(rr[i].size.width, rr[j].size.width);
            if (contour_len1 > 0.7 || contour_len2 > 0.7)
                continue;

            RotatedRect Armor;
            Armor.center.x = (rr[i].center.x + rr[j].center.x) / 2;
            Armor.center.y = (rr[i].center.y + rr[j].center.y) / 2;
            float rh, rw, y_diff, x_diff;
            rh = (rr[i].size.height + rr[j].size.height) / 2;
            rw = sqrt((rr[i].center.x - rr[j].center.x) * (rr[i].center.x - rr[j].center.x) + (rr[i].center.y - rr[j].center.y) * (rr[i].center.y - rr[j].center.y));
            float ratio = rw / rh;
            x_diff = abs(rr[i].center.x - rr[j].center.x) / rh;
            y_diff = abs(rr[i].center.y - rr[j].center.y) / rh;
            if (ratio < 0.6 || ratio > 5.0 || x_diff < 0.5 || y_diff > 2.0)
                continue;

            Armor.size.height = rh;
            Armor.size.width = rw;
            vRec.push_back(Armor);

            Point2f P[4];
            Armor.points(P);
            for (int j = 0; j < 4; ++j)
            {
                line(dst, P[j], P[(j + 2) % 4], Scalar(0, 255, 0), 3);
            }
            circle(dst, Armor.center, 3, CV_RGB(255, 0, 0), -1);

            Point2f P1[4];
            rr[i].points(P1);
            points.clear();
            points.push_back(P1[1]);
            points.push_back(P1[2]);
            points.push_back(P1[3]);
            points.push_back(P1[4]);

            solvePnP(armor, points, camera_matrix, distortion_coefficients, r_vec, t_vec, SOLVEPNP_ITERATIVE);
            cout << "t_vec=" << t_vec << endl;
            cout << "distant=" << t_vec.row(2) <<endl;
            
        }
    }
}


int main()
{

    Mat frame;
    VideoCapture capture("2.avi");
    if (!capture.isOpened())
    {
        cout << " cannot open the vedio" << endl;
        return -1;
    }

    capture >> frame;
    while (capture.read(frame))
    {

        Mat hsvimg;
        cvtColor(frame, hsvimg, COLOR_BGR2HSV);
        blur(hsvimg, hsvimg, Size(3, 3));
        vector<Mat> hsvsplit;
        split(hsvimg, hsvsplit);
        equalizeHist(hsvsplit[2], hsvsplit[2]);
        merge(hsvsplit, hsvimg);

        Mat threshold_img;
        inRange(hsvimg, Scalar(low_h, low_s, low_v), Scalar(max_h, max_s, max_v), threshold_img);
        Canny(threshold_img, threshold_img, 3, 9, 3);
        Mat element = getStructuringElement(MORPH_RECT, Size(3, 3));
        dilate(threshold_img, threshold_img, element);

        draw(threshold_img, frame);

        imshow("AD", frame);
        if (waitKey(30) == 32)
        {
            if (waitKey(0) == 32)
            {
                break;
            }
        }
    }
    capture.release();
    return 0;
}
